# Copyright (c) Alibaba, Inc. and its affiliates.
from dataclasses import dataclass, field
from typing import List, Optional

from swift.utils import get_logger

logger = get_logger()


@dataclass
class GenerationArguments:
    """
    GenerationArguments 类是一个数据类，用于封装文本生成相关的各种参数。

    参数说明：
        max_new_tokens (Optional[int]): 生成的最大新 token 数，默认无限制。
        temperature (Optional[float]): 控制采样的随机性，值越低越确定性，默认 None。
        top_k (Optional[int]): top-k 采样参数，仅采样概率前 k 个 token，默认 None。
        top_p (Optional[float]): nucleus sampling（累积概率）参数，默认 None。
        repetition_penalty (Optional[float]): 控制重复生成的惩罚项，默认 None。
        num_beams (int): beam search 的宽度（分支数），默认 1，表示不开启。
        stream (bool): 是否启用流式生成（边生成边输出），默认 False。
        stop_words (List[str]): 停止生成的词汇列表，遇到即停止生成，默认空列表。
        logprobs (bool): 是否返回 token 的概率信息。
        top_logprobs (Optional[int]): 返回每个 token 最可能的前 top_logprobs 个概率。
    """

    # 生成参数配置

    max_new_tokens: Optional[int] = None  # 限制生成的新 token 数量，None 表示不限制。
    temperature: Optional[float] = None  # 温度系数，控制生成随机性。
    top_k: Optional[int] = None  # top-k 采样策略参数。
    top_p: Optional[float] = None  # top-p（nucleus）采样策略参数。
    repetition_penalty: Optional[float] = None  # 对重复内容的惩罚。
    num_beams: int = 1  # beam search 分支数量，默认关闭（设置为 1）。

    stream: bool = False  # 是否启用流式输出。
    stop_words: List[str] = field(default_factory=list)  # 停止词列表，遇到即中断生成。
    logprobs: bool = False  # 是否输出每个 token 的对数概率。
    top_logprobs: Optional[int] = None  # 每个 token 输出 top N 的概率。

    # 获取生成请求配置的方法
    def get_request_config(self):
        # 如果任务类型不是 causal language model，则直接返回 None
        if getattr(self, 'task_type') != 'causal_lm':
            return

        # 引入 RequestConfig 配置类（用于封装生成请求）
        from swift.llm import RequestConfig

        # 返回一个 RequestConfig 实例，包含当前对象的生成参数
        return RequestConfig(
            max_tokens=self.max_new_tokens,               # 最大生成 token 数
            temperature=self.temperature,                 # 采样温度
            top_p=self.top_p,                             # top-p 采样
            top_k=self.top_k,                             # top-k 采样
            num_beams=self.num_beams,                     # beam search 分支数
            stop=self.stop_words,                         # 停止词
            stream=self.stream,                           # 是否流式输出
            repetition_penalty=self.repetition_penalty,   # 重复惩罚
            logprobs=self.logprobs,                       # 是否输出 logprob
            top_logprobs=self.top_logprobs)               # 输出 top logprobs 的数量
